import asyncio
import websockets
from config.logger import setup_logging

TAG = __name__
logger = setup_logging()

class WSClient:
    def __init__(self, uri):
        self.uri = uri
        self.websocket = None
        self.connected = False
        self.connection_lock = asyncio.Lock()
        self.wsServer = None

    async def connect(self, server=None ):
        async with self.connection_lock:
            if self.connected:
                return
            
            try:
                self.wsServer = server
                self.websocket = await websockets.connect(self.uri)
                self.connected = True
                logger.bind(tag=TAG).info(f"Connected to {self.uri}")
                asyncio.create_task(self._heartbeat())
            except Exception as e:
                logger.bind(tag=TAG).error(f"Connection failed: {e}")
                raise

    async def _heartbeat(self):
        while self.connected:
            try:
                await self.websocket.ping()
                await asyncio.sleep(30)  # Send ping every 30 seconds
            except:
                self.connected = False
                logger.bind(tag=TAG).info(f"Connection lost, attempting to reconnect...")
                await self._reconnect()

    async def _reconnect(self):
        try:
            await self.connect()
        except Exception as e:
            logger.bind(tag=TAG).info(f"Reconnection failed: {e}")
            await asyncio.sleep(5)  # Wait before next reconnection attempt
            await self._reconnect()

    async def send_message(self, message):
        if not self.connected:
            raise Exception("WebSocket connection not established")
            
        try:
            await self.websocket.send(message)
            # logger.bind(tag=TAG).info(f"Sent: {message}")
            # logger.bind(tag=TAG).info(f"Sent: Msg")
        except Exception as e:
            self.connected = False
            logger.bind(tag=TAG).error(f"Failed to send message: {e}")
            raise

    async def receive_message(self):
        if not self.connected:
            raise Exception("WebSocket connection not established")
            
        try:
            message = await self.websocket.recv()
            logger.bind(tag=TAG).info(f"Received: {message}")

            # self.wsServer.broadcast_message("test")

            return message
        except Exception as e:
            self.connected = False
            logger.bind(tag=TAG).error(f"Failed to receive message: {e}")
            raise

    async def close(self):
        async with self.connection_lock:
            if self.websocket and self.connected:
                try:
                    await self.websocket.close()
                    self.connected = False
                    self.logger.info("WebSocket connection closed")
                except Exception as e:
                    self.logger.error(f"Error closing connection: {e}")

# Singleton instance for global access
ws_client = WSClient("ws://127.0.0.1:14167/echo?id=esp_server")